
#include "StdAfx.h"

#include <windows.h>

#include "CSymbolFinder.h"

BOOL CSymbolFinder::Init(const TCHAR* pszPdbFile)
{
	m_wstrPdbFile=pszPdbFile;

	IDiaDataSourcePtr pDiaSrc(CLSID_DiaSource);	// note to call CoInitializeEx, all exception will generate
	m_pDiaSrc=pDiaSrc;
	if(FAILED(m_pDiaSrc->loadDataFromPdb(m_wstrPdbFile.c_str())))
	{
		return FALSE;
	}

	if(FAILED( m_pDiaSrc->openSession( &m_pSession )))
	{
		return FALSE;
	}

	if(FAILED( m_pSession->get_globalScope( &m_pGlobal )))
	{
		return FALSE;
	}
	
	return TRUE;
};


BOOL CSymbolFinder::FindSymbolS(const TCHAR* pszName, vector<IDiaSymbolPtr>* pv, enum SymTagEnum symtag=SymTagNull)
{
	if(0==ForEachMatchSymblol(OnMatchFindSymbolS<IDiaSymbolPtr>, pszName, symtag))
	{
		return FALSE;
	}
	
	return TRUE;
};

DWORD CSymbolFinder::ForEachMatchSymblol(pfnOnMatchSymbol pfn, VOID* pv, const TCHAR* pszName, enum SymTagEnum symtag)
{
	enum SymTagEnum enTpS[]={SymTagUDT, SymTagTypedef, SymTagEnum, SymTagBaseType, SymTagData};
	DWORD dwDim=ARRAYOF(enTpS);
	DWORD dwRet=0;

	if(symtag!=SymTagNull)
	{
		enTpS[0]=symtag;
		dwDim=1;
	}

	for(DWORD i=0; i<dwDim; i++)
	{
		IDiaEnumSymbolsPtr pEnum;
		if( (SUCCEEDED(m_pGlobal->findChildren(enTpS[i], pszName, m_dwCompareFlags, &pEnum)))
			&& (pEnum!=NULL) ) // sometimes even findchidren success, but the pEnum is null!!!
		{
			IDiaSymbolPtr psymbolChild;
			DWORD celt;
			for( DWORD dwChldCnt=0;
				SUCCEEDED( pEnum->Next( 1, &psymbolChild, &celt ) ) && celt == 1;
				dwChldCnt++) 
			{
                        dwRet++;
                        if( (this->*pfn)(pv, psymbolChild, enTpS[i]) )
                        {
                            return dwRet;
                        }
				psymbolChild=NULL;
			}
		
		}
	}
	
	return dwRet;
}

template<IDiaSymbol*, BOOL RET>
BOOL CSymbolFinder::OnMatchSymbol(VOID* pv, IDiaSymbol* psymb, enum SymTagEnum symtag)
{
    vector<IDiaSymbol*>* pvt=(vector<IDiaSymbol*>*)pv;
    psymb->AddRef();
    pvt->push_back(psymb);
    return RET;
}

template<IDiaSymbolPtr, BOOL RET>
BOOL CSymbolFinder::OnMatchSymbol(VOID* pv, IDiaSymbol* psymb, enum SymTagEnum symtag)
{
    vector<IDiaSymbolPtr>* pvt=(vector<IDiaSymbolPtr>*)pv;
    pvt->push_back(psymb);
    return RET;
}

BOOL CSymbolFinder::FindBasicSymbol(BasicType btBaseType, DWORD dwLen, IDiaSymbol** psymb)
{
	IDiaEnumSymbolsPtr pEnum;
	if( (FAILED(m_pGlobal->findChildren(SymTagBaseType, NULL, 0, &pEnum)))
		|| (pEnum==NULL) ) // sometimes even findchidren success, but the pEnum is null!!!
	{
		return NULL;
	}
	
	IDiaSymbolPtr psymbolChild;
	DWORD celt;
	for( DWORD dwChldCnt=0;
		SUCCEEDED( pEnum->Next( 1, &psymbolChild, &celt ) ) && celt == 1;
		dwChldCnt++) 
	{
		ULONGLONG ull;
		VERIFY(SUCCEEDED(psymbolChild->get_length(&ull)));
		if(dwLen==(DWORD)(ull))
		{
			*psymb=psymbolChild;
			(*psymb)->AddRef();
			return TRUE;
		}
		psymbolChild=NULL;
	}
	
	return FALSE;
};
